import numpy as np
import pandas as pd
import os
import torch
from torch.utils.data import Dataset

def load_data(subject_dir, csv_path):
    df = pd.read_csv(csv_path, index_col=0)
    subjects = os.listdir(subject_dir)

    x = []
    y = []
    for subject in subjects:
        features_path = os.path.join(subject_dir, subject)
        if not os.path.exists(features_path) or not features_path.endswith('npy'):
            continue
        else:
            row = df.loc[subject.split('.')[0]]
            label = int(row['Label'])

            x.append(np.load(features_path))
            y.append(label)

    x = np.array(x)
    y = np.array(y)
    return x, y

class AD_Dataset(Dataset):
    def __init__(self, x, y, index,device):
        # 数据预处理
        x = x[index]
        y = y[index]

        self.x = torch.from_numpy(x).to(torch.float32)
        self.y = torch.from_numpy(y).long()
        self.index = index
        self.device = device
        label_count_dict={}
        for label in y:
            label_count_dict[label] = label_count_dict.get(label,0)+1
        print(label_count_dict)

        self.class_weight = []
        for i, (cla, count) in enumerate(sorted(label_count_dict.items(), key=lambda x: x[0])):
            weight = 1. / (count + 1)
            self.class_weight.append(weight)
        self.sample_weight = [self.class_weight[label] for label in y]

    def __getitem__(self, index):
        xi = self.x[index]
        yi = self.y[index]
        i = self.index[index]
        return xi, yi,i

    def __len__(self):
        return len(self.y)

if __name__=='__main__':
    ds = AD_Dataset(r'D:\dataset\AD\train_data\train',r'D:\dataset\AD\train_data\train_open.csv')
    print(ds.__len__())
    print(ds.__getitem__(0))